package ch.ethz.fcl.input;

import java.nio.ShortBuffer;
import java.util.ArrayList;
import java.util.List;

import org.OpenNI.Context;
import org.OpenNI.DepthGenerator;
import org.OpenNI.DepthMetaData;
import org.OpenNI.GeneralException;
import org.OpenNI.OutArg;
import org.OpenNI.ScriptNode;
import org.opencv.core.CvType;
import org.opencv.core.Mat;

import ch.ethz.fcl.Constant;
import ch.ethz.fcl.KinectContainer;

/**
 * 16-bit depth value retrieved from Kinect
 * 
 * @author crazyfrog
 * 
 */

public class DepthImage {
	private KinectContainer container;

	private static final int ACCUMULATE = 5;

	private OutArg<ScriptNode> scriptNode;
	private Context context;
	public DepthGenerator depthGen;

	private Mat[] imageMats;
	private short[][] imgData;

	private short[] allLayersImgData;
	private Mat allLayersImageMat;

	private DepthMetaData depthMD;
	private List<short[]> histDepthValue;

	private final int LAYERS = Constant.HEIGHTS.length;

	public DepthImage(KinectContainer container) {
		this.container = container;
		init();
	}

	public DepthImage(KinectContainer container, String xml_config_file) {
		this.container = container;
		try {
			scriptNode = new OutArg<ScriptNode>();
			context = Context.createFromXmlFile(xml_config_file, scriptNode);

			depthGen = DepthGenerator.create(context);

			depthMD = depthGen.getMetaData();
		} catch (GeneralException e) {
			e.printStackTrace();
			System.exit(1);
		}

		init();
	}

	private void init() {
		imgData = new short[LAYERS][Coordinate.WIDTH * Coordinate.HEIGHT];
		histDepthValue = new ArrayList<short[]>();

		imageMats = new Mat[LAYERS];
		for (int i = 0; i < LAYERS; i++)
			imageMats[i] = createImageMat(imgData[i], Coordinate.WIDTH,
					Coordinate.HEIGHT);

		allLayersImgData = new short[Coordinate.WIDTH * Coordinate.HEIGHT];
		allLayersImageMat = createImageMat(allLayersImgData, Coordinate.WIDTH,
				Coordinate.HEIGHT);
	}

	public void update() {
		try {
			depthMD = depthGen.getMetaData();
			context.waitAnyUpdateAll();

			saveHistDepthValue(depthMD);

			// for (int i = 0; i < Coordinate.HEIGHT; i++)
			// for (int j = 0; j < Coordinate.WIDTH; j++)
			// allLayersImgData[i * Coordinate.WIDTH + j] = (short) getDepth(j,
			// i);
			updateFromDepthMap(container.getTableDepthMap().getHistDepthMap(),
					container.getTableDepthMap().getHeightMap());
		} catch (GeneralException e) {
			e.printStackTrace();
		}
	}

	public void updateFromDepthMap(short[][] depthMap, short[][] heightMap) {
		int scale = 20;
		allLayersImageMat = Mat.zeros(Coordinate.HEIGHT, Coordinate.WIDTH,
				CvType.CV_16U);
		for (int layer = 0; layer < LAYERS; layer++)
			imageMats[layer] = Mat.zeros(Coordinate.HEIGHT, Coordinate.WIDTH,
					CvType.CV_16U);

		for (int row = 0; row < Coordinate.HEIGHT; row++)
			for (int col = 0; col < Coordinate.WIDTH; col++) {
				short pixelHeight = heightMap[row][col];

				if (pixelHeight < 200)
					allLayersImageMat.put(row, col, depthMap[row][col] * scale);

				if (pixelHeight == 0) {
					continue;
				}

				int heightLayer = -1;
				for (int layer = 0; layer < LAYERS; layer++) {
					int layerHeight = Constant.HEIGHTS[layer];

					if (pixelHeight > layerHeight - Constant.HEIGHT_VARIANCE
							&& pixelHeight < layerHeight
									+ Constant.HEIGHT_VARIANCE) {
						heightLayer = layer;
						break;
					}
				}
				if (heightLayer != -1)
					imageMats[heightLayer].put(row, col, depthMap[row][col]
							* scale);
			}
	}

	private void saveHistDepthValue(DepthMetaData md) {
		if (histDepthValue.size() == ACCUMULATE)
			histDepthValue.remove(0);

		short[] depthVal = new short[Coordinate.WIDTH * Coordinate.HEIGHT];
		ShortBuffer buf = md.getData().createShortBuffer();
		buf.rewind();
		for (int j = 0; j < Coordinate.WIDTH * Coordinate.HEIGHT; j++)
			depthVal[j] = buf.get(j);
		histDepthValue.add(depthVal);
	}

	private Mat createImageMat(short[] data, int width, int height) {
		int scale = 20;
		Mat mat = new Mat(height, width, CvType.CV_16U);
		for (int row = 0; row < height; row++)
			for (int col = 0; col < width; col++)
				mat.put(row, col, data[row * width + col] * scale);

		return mat;
	}

	/**
	 * get the depth at position (x, y)
	 * 
	 * @param x
	 * @param y
	 * @return
	 */
	public float getDepth(int x, int y) {
		float depth = 0;
		int pos = x + y * Coordinate.WIDTH;
		for (int i = 0; i < histDepthValue.size(); i++)
			depth += histDepthValue.get(i)[pos];

		return depth / histDepthValue.size();
	}

	public DepthMetaData getDepthMD() {
		return depthMD;
	}

	public Mat[] getImageMats() {
		return imageMats;
	}

	public Mat getAllLayerImageMat() {
		return allLayersImageMat;
	}
}
